import copy
import math
from pathlib import Path
import os
import logging
from scipy import ndimage
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

ATTENTION_Q = "MultiHeadDotProductAttention_1/query"
ATTENTION_K = "MultiHeadDotProductAttention_1/key"
ATTENTION_V = "MultiHeadDotProductAttention_1/value"
ATTENTION_OUT = "MultiHeadDotProductAttention_1/out"
FC_0 = "MlpBlock_3/Dense_0"
FC_1 = "MlpBlock_3/Dense_1"
ATTENTION_NORM = "LayerNorm_0"
MLP_NORM = "LayerNorm_2"


def np2th(weights, conv=False):
    """Possibly convert HWIO to OIHW."""
    if conv:
        weights = weights.transpose([3, 2, 0, 1])
    return torch.from_numpy(weights)


def swish(x):
    return x * torch.sigmoid(x)


ACTION2FN = {"gelu": torch.nn.GELU(),
             "relu": torch.nn.ReLU(),
             "swish": swish}


class Embeddings(nn.Module):
    """
    Construct the embeddings from patch, position embeddings.
    """

    def __init__(self,
                 in_channel=3,
                 img_size=224,
                 patch_size=16,
                 slide_step=12,
                 split_type='non-overlap',
                 hidden_size=768,
                 dropout_rate=0.1,
                 ):
        super(Embeddings, self).__init__()
        if isinstance(img_size, int):
            img_size = (img_size, img_size)
        if not isinstance(img_size, tuple):
            img_size = tuple(img_size)

        self.hybrid = None

        if split_type == 'non-overlap':
            n_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size)
            self.patch_embeddings = nn.Conv2d(in_channel, hidden_size, kernel_size=patch_size, stride=patch_size)
        elif split_type == 'overlap':
            n_patches = ((img_size[0] - patch_size) // slide_step + 1) * ((img_size[1] - patch_size) // slide_step + 1)
            self.patch_embeddings = nn.Conv2d(in_channel, hidden_size, kernel_size=patch_size, stride=slide_step)
        self.position_embeddings = nn.Parameter(torch.zeros(1, n_patches + 1, hidden_size))
        self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size))
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x):
        B, C, H, W = x.shape
        cls_tokens = self.cls_token.expand(B, -1, -1)

        if self.hybrid:
            x = self.hybrid_model(x)
        x = self.patch_embeddings(x)  # b,c,h,w
        x = x.flatten(2)  # b,c,h,w -> b,c h*w
        x = x.transpose(-1, -2)  # b,c,h*w -> b,h*w,c
        x = torch.cat((cls_tokens, x), dim=1)  # b,h*w+1,c
        embeddings = x + self.position_embeddings
        embeddings = self.dropout(embeddings)
        return embeddings


class Encoder(nn.Module):
    def __init__(self,
                 num_layers=12,
                 hidden_size=768,
                 num_attention_heads=12,
                 mlp_dim=3072,
                 action='gelu',
                 dropout_rate=0.1,
                 attention_dropout_rate=0,
                 proj_dropout_rate=0,
                 ):
        super(Encoder, self).__init__()
        self.layer = nn.ModuleList()
        for _ in range(num_layers - 1):
            layer = Block(hidden_size=hidden_size, mlp_dim=mlp_dim, num_attention_heads=num_attention_heads,
                          action=action,
                          dropout_rate=dropout_rate, attention_dropout_rate=attention_dropout_rate,
                          proj_dropout_rate=proj_dropout_rate)
            self.layer.append(copy.deepcopy(layer))
        self.part_select = Part_Attention()
        self.part_layer = Block(hidden_size=hidden_size, mlp_dim=mlp_dim, num_attention_heads=num_attention_heads,
                                action=action,
                                dropout_rate=dropout_rate, attention_dropout_rate=attention_dropout_rate,
                                proj_dropout_rate=proj_dropout_rate)
        self.part_norm = nn.LayerNorm(hidden_size, eps=1e-6)

    def forward(self, hidden_states):
        attn_weights = []
        for layer in self.layer:
            hidden_states, weights = layer(hidden_states)
            attn_weights.append(weights)
        part_num, part_inx = self.part_select(attn_weights)
        part_inx = part_inx + 1
        parts = []
        B, num = part_inx.shape
        for i in range(B):
            parts.append(hidden_states[i, part_inx[i, :]])
        parts = torch.stack(parts).squeeze(1)
        concat = torch.cat((hidden_states[:, 0].unsqueeze(1), parts), dim=1)
        part_states, part_weights = self.part_layer(concat)
        part_encoded = self.part_norm(part_states)

        return part_encoded


class Part_Attention(nn.Module):
    def __init__(self):
        super(Part_Attention, self).__init__()

    def forward(self, x):
        length = len(x)
        last_map = x[0]
        for i in range(1, length):
            last_map = torch.matmul(x[i], last_map)
        last_map = last_map[:, :, 0, 1:]
        _, max_inx = last_map.max(2)
        return _, max_inx


class Attention(nn.Module):
    def __init__(self,
                 hidden_size=768,
                 num_attention_heads=12,
                 attention_dropout_rate=0.0,
                 proj_dropout_rate=0.0
                 ):
        super(Attention, self).__init__()
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.attention_head_size = int(hidden_size / num_attention_heads)
        self.all_head_size = num_attention_heads * self.attention_head_size

        self.query = nn.Linear(hidden_size, self.all_head_size)
        self.key = nn.Linear(hidden_size, self.all_head_size)
        self.value = nn.Linear(hidden_size, self.all_head_size)

        self.out = nn.Linear(hidden_size, hidden_size)
        self.attn_dropout = nn.Dropout(attention_dropout_rate)
        self.proj_dropout = nn.Dropout(proj_dropout_rate)

        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        mixed_query_layer = self.query(x)
        mixed_key_layer = self.key(x)
        mixed_value_layer = self.value(x)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        attention_scores = query_layer @ key_layer.transpose(-1, -2)
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        attention_probs = self.softmax(attention_scores)
        weights = attention_probs
        attention_probs = self.attn_dropout(attention_probs)

        context_layer = attention_probs @ value_layer
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)
        attention_output = self.out(context_layer)
        attention_output = self.proj_dropout(attention_output)
        return attention_output, weights

    def transpose_for_scores(self, x):
        # B L C -> B L HEAD HEAD_SIZE
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)  # B HEAD L HEAD_SIZE


class Block(nn.Module):
    def __init__(self,
                 hidden_size,
                 mlp_dim,
                 num_attention_heads=12,
                 attention_norm=None,
                 ffn_norm=None,
                 action='gelu',
                 dropout_rate=0.1,
                 attention_dropout_rate=0,
                 proj_dropout_rate=0):
        super(Block, self).__init__()
        if attention_norm is None:
            attention_norm = nn.LayerNorm
        if ffn_norm is None:
            ffn_norm = nn.LayerNorm

        self.hidden_size = hidden_size
        self.attention_norm = attention_norm(self.hidden_size, eps=1e-6)
        self.ffn_norm = ffn_norm(self.hidden_size, eps=1e-6)
        self.ffn = MLP(hidden_size=hidden_size, mlp_dim=mlp_dim, action=action, dropout_rate=dropout_rate)
        self.attn = Attention(hidden_size=hidden_size, num_attention_heads=num_attention_heads,
                              attention_dropout_rate=attention_dropout_rate,
                              proj_dropout_rate=proj_dropout_rate)

    def forward(self, x):
        h = x
        x = self.attention_norm(x)
        x, weights = self.attn(x)
        x = x + h

        h = x
        x = self.ffn_norm(x)
        x = self.ffn(x)
        x = x + h
        return x, weights

    def load_from(self, weights, n_block):
        ROOT = f"Transformer/encoderblock_{n_block}"
        with torch.no_grad():
            query_weight = np2th(weights[ROOT + "/" + ATTENTION_Q + "/kernel"]).view(self.hidden_size,
                                                                                     self.hidden_size).t()
            key_weight = np2th(weights[ROOT + "/" + ATTENTION_K + "/kernel"]).view(self.hidden_size,
                                                                                   self.hidden_size).t()
            value_weight = np2th(weights[ROOT + "/" + ATTENTION_V + "/kernel"]).view(self.hidden_size,
                                                                                     self.hidden_size).t()
            out_weight = np2th(weights[ROOT + "/" + ATTENTION_OUT + "/kernel"]).view(self.hidden_size,
                                                                                     self.hidden_size).t()

            query_bias = np2th(weights[ROOT + "/" + ATTENTION_Q + "/bias"]).view(-1)
            key_bias = np2th(weights[ROOT + "/" + ATTENTION_K + "/bias"]).view(-1)
            value_bias = np2th(weights[ROOT + "/" + ATTENTION_V + "/bias"]).view(-1)
            out_bias = np2th(weights[ROOT + "/" + ATTENTION_OUT + "/bias"]).view(-1)

            self.attn.query.weight.copy_(query_weight)
            self.attn.key.weight.copy_(key_weight)
            self.attn.value.weight.copy_(value_weight)
            self.attn.out.weight.copy_(out_weight)
            self.attn.query.bias.copy_(query_bias)
            self.attn.key.bias.copy_(key_bias)
            self.attn.value.bias.copy_(value_bias)
            self.attn.out.bias.copy_(out_bias)

            mlp_weight_0 = np2th(weights[ROOT + "/" + FC_0 + "/kernel"]).t()
            mlp_weight_1 = np2th(weights[ROOT + "/" + FC_1 + "/kernel"]).t()
            mlp_bias_0 = np2th(weights[ROOT + "/" + FC_0 + "/bias"]).t()
            mlp_bias_1 = np2th(weights[ROOT + "/" + FC_1 + "/bias"]).t()

            self.ffn.fc1.weight.copy_(mlp_weight_0)
            self.ffn.fc2.weight.copy_(mlp_weight_1)
            self.ffn.fc1.bias.copy_(mlp_bias_0)
            self.ffn.fc2.bias.copy_(mlp_bias_1)

            self.attention_norm.weight.copy_(np2th(weights[ROOT + "/" + ATTENTION_NORM + "/scale"]))
            self.attention_norm.bias.copy_(np2th(weights[ROOT + "/" + ATTENTION_NORM + "/bias"]))
            self.ffn_norm.weight.copy_(np2th(weights[ROOT + "/" + MLP_NORM + "/scale"]))
            self.ffn_norm.bias.copy_(np2th(weights[ROOT + "/" + MLP_NORM + "/bias"]))


class MLP(nn.Module):
    def __init__(self, hidden_size, mlp_dim, action='gelu', dropout_rate=0.1):
        super(MLP, self).__init__()

        assert action in ACTION2FN.keys(), f"action function must is {ACTION2FN.keys()}, now action is {action}"

        self.fc1 = nn.Linear(hidden_size, mlp_dim)
        self.fc2 = nn.Linear(mlp_dim, hidden_size)
        self.act_fn = ACTION2FN[action]
        self.dropout = nn.Dropout(dropout_rate)
        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.fc1.weight)
        nn.init.xavier_uniform_(self.fc2.weight)
        nn.init.normal_(self.fc1.bias, std=1e-6)
        nn.init.normal_(self.fc2.bias, std=1e-6)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act_fn(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x


class Transformer(nn.Module):
    def __init__(self, cfg='config/example.yaml'):
        super(Transformer, self).__init__()
        if isinstance(cfg, dict):
            config = cfg
        else:
            import yaml
            file_name = Path(cfg).name
            with open(cfg, encoding='utf-8') as f:
                config = yaml.safe_load(f)

        img_size = config["model"]["image_size"]
        patch_size = config["model"]["patches"]["patch_size"]
        split_type = config["model"]["patches"]["split_type"]
        hidden_size = config["model"]["patches"]["hidden_size"]
        slide_step = config["model"]["patches"]["slide_step"]
        dropout_rate = config["model"]["transformer"]["dropout_rate"]
        num_layers = config["model"]["transformer"]["num_layers"]
        mlp_dim = config["model"]["transformer"]["mlp_dim"]
        action = config["model"]["transformer"]["action"]
        num_attention_heads = config["model"]["transformer"]["num_heads"]
        attention_dropout_rate = config["model"]["transformer"]["attention_dropout_rate"]
        proj_dropout_rate = config["model"]["transformer"]["attention_dropout_rate"]

        self.embeddings = Embeddings(in_channel=3, img_size=img_size, patch_size=patch_size, slide_step=slide_step,
                                     split_type=split_type, hidden_size=hidden_size, dropout_rate=dropout_rate)
        self.encoder = Encoder(num_layers=num_layers, hidden_size=hidden_size, num_attention_heads=num_attention_heads,
                               mlp_dim=mlp_dim, action=action,
                               dropout_rate=dropout_rate, attention_dropout_rate=attention_dropout_rate,
                               proj_dropout_rate=proj_dropout_rate)

    def forward(self, x):
        embedding_out = self.embeddings(x)
        part_encoded = self.encoder(embedding_out)
        return part_encoded


class VisionTransformer(nn.Module):
    def __init__(self, cfg='config/example.yaml', num_classes=21843, smoothing_value=0., zero_head=False):
        super(VisionTransformer, self).__init__()
        if isinstance(cfg, dict):
            config = cfg
        else:
            import yaml
            file_name = Path(cfg).name
            with open(cfg, encoding='utf-8') as f:
                config = yaml.safe_load(f)

        self.num_classes = num_classes
        self.smoothing_value = smoothing_value
        self.zero_head = zero_head
        self.classifier = config["model"]["classifier"]
        self.part_head = nn.Linear(config["model"]["patches"]["hidden_size"], num_classes)

        self.transformer = Transformer(config)

    def forward(self, x, label=None):
        part_tokens = self.transformer(x)
        part_logits = self.part_head(part_tokens[:, 0])

        if label is not None:
            if self.smoothing_value == 0:
                loss_fct = nn.CrossEntropyLoss()
            else:
                from losses.labelSmoothing import LabelSmoothing
                loss_fct = LabelSmoothing(self.smoothing_value)
            part_loss = loss_fct(part_logits.view(-1, self.num_classes), label.view(-1))
            from losses.contrastive_loss import contrastive_loss
            contrast_loss = contrastive_loss(part_tokens[:, 0], label.view(-1))
            loss = part_loss + contrast_loss
            return loss, part_logits
        else:
            return part_logits

    def load_from(self, weights):
        with torch.no_grad():
            self.transformer.embeddings.patch_embeddings.weight.copy_(np2th(weights["embedding/kernel"], conv=True))
            self.transformer.embeddings.patch_embeddings.bias.copy_(np2th(weights["embedding/bias"]))
            self.transformer.embeddings.cls_token.copy_(np2th(weights["cls"]))
            self.transformer.encoder.part_norm.weight.copy_(np2th(weights["Transformer/encoder_norm/scale"]))
            self.transformer.encoder.part_norm.bias.copy_(np2th(weights["Transformer/encoder_norm/bias"]))

            posemb = np2th(weights["Transformer/posembed_input/pos_embedding"])
            posemb_new = self.transformer.embeddings.position_embeddings
            if posemb.size() == posemb_new.size():
                self.transformer.embeddings.position_embeddings.copy_(posemb)
            else:
                logging.info("load_pretrained: resized variant: %s to %s" % (posemb.size(), posemb_new.size()))
                ntok_new = posemb_new.size(1)

                if self.classifier == "token":
                    posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
                    ntok_new -= 1
                else:
                    posemb_tok, posemb_grid = posemb[:, :0], posemb[0]

                gs_old = int(np.sqrt(len(posemb_grid)))
                gs_new = int(np.sqrt(ntok_new))
                print('load_pretrained: grid-size from %s to %s' % (gs_old, gs_new))
                posemb_grid = posemb_grid.reshape(gs_old, gs_old, -1)

                zoom = (gs_new / gs_old, gs_new / gs_old, 1)
                posemb_grid = ndimage.zoom(posemb_grid, zoom, order=1)
                posemb_grid = posemb_grid.reshape(1, gs_new * gs_new, -1)
                posemb = np.concatenate([posemb_tok, posemb_grid], axis=1)
                self.transformer.embeddings.position_embeddings.copy_(np2th(posemb))

            for bname, block in self.transformer.encoder.named_children():
                if bname.startswith('part') == False:
                    for uname, unit in block.named_children():
                        unit.load_from(weights, n_block=uname)

            if self.transformer.embeddings.hybrid:
                self.transformer.embeddings.hybrid_model.root.conv.weight.copy_(
                    np2th(weights["conv_root/kernel"], conv=True))
                gn_weight = np2th(weights["gn_root/scale"]).view(-1)
                gn_bias = np2th(weights["gn_root/bias"]).view(-1)
                self.transformer.embeddings.hybrid_model.root.gn.weight.copy_(gn_weight)
                self.transformer.embeddings.hybrid_model.root.gn.bias.copy_(gn_bias)

                for bname, block in self.transformer.embeddings.hybrid_model.body.named_children():
                    for uname, unit in block.named_children():
                        unit.load_from(weights, n_block=bname, n_unit=uname)


if __name__ == '__main__':
    x = torch.randn(1, 3, 448, 448)
    model = VisionTransformer(cfg='../config/example.yaml')
    model.load_from(np.load('../pretrain_weights/imagenet21k_ViT-B_16.npz'))
    y = model(x)
    print(model)
    print(y.shape)
